package handlers

import (
	"context"
	"crypto/sha1"
	"crypto/x509"
	"encoding/asn1"
	"encoding/base64"
	"encoding/xml"
	"fmt"
	"log/slog"
	"net/http"
	"regexp"
	"strings"

	"github.com/beevik/etree"
	"github.com/gin-gonic/gin"
	"github.com/russellhaering/goxmldsig"

	"chinaport-proxy/protocol"
	"chinaport-proxy/utils"
)

type CEBSignRequest struct {
	XML string `json:"xml" binding:"required"`
}

type CEBSignResponse struct {
	SignedXML string `json:"signedXml"`
	Success   bool   `json:"success"`
	Error     string `json:"error,omitempty"`
}

type XMLSignature struct {
	XMLName        xml.Name   `xml:"ds:Signature"`
	XmlnsDS        string     `xml:"xmlns:ds,attr"`
	SignedInfo     SignedInfo `xml:"ds:SignedInfo"`
	SignatureValue string     `xml:"ds:SignatureValue"`
	KeyInfo        KeyInfo    `xml:"ds:KeyInfo"`
}

type KeyInfo struct {
	XMLName  xml.Name `xml:"ds:KeyInfo"`
	KeyName  string   `xml:"ds:KeyName"`
	X509Data X509Data `xml:"ds:X509Data"`
}

type X509Data struct {
	XMLName         xml.Name `xml:"ds:X509Data"`
	X509Certificate string   `xml:"ds:X509Certificate"`
}

type SignedInfo struct {
	XMLName                xml.Name  `xml:"ds:SignedInfo"`
	CanonicalizationMethod Method    `xml:"ds:CanonicalizationMethod"`
	SignatureMethod        Method    `xml:"ds:SignatureMethod"`
	Reference              Reference `xml:"ds:Reference"`
}

type Reference struct {
	XMLName      xml.Name   `xml:"ds:Reference"`
	URI          string     `xml:"URI,attr"`
	Transforms   Transforms `xml:"ds:Transforms"`
	DigestMethod Method     `xml:"ds:DigestMethod"`
	DigestValue  string     `xml:"ds:DigestValue"`
}

type Transforms struct {
	Transform []Method `xml:"ds:Transform"`
}

type Method struct {
	Algorithm string `xml:"Algorithm,attr"`
}

type CertificateInfo struct {
	KeyName       string
	Certificate   string
	CertificateNo string
}

// determineSignatureMethod attempts to determine the signature algorithm to advertise in <ds:SignatureMethod>.
// Defaults to RSA-SHA1. If the certificate appears to be non-RSA (e.g., SM2/EC), returns SM2-SM3 URI.
func determineSignatureMethod(base64Cert string) string {
	// Default RSA-SHA1
	const (
		rsaSHA1 = "http://www.w3.org/2000/09/xmldsig#rsa-sha1"
		sm2SM3  = "http://www.w3.org/2000/09/xmldsig#sm2-sm3"
	)
	// Decode base64 DER
	certBody := normalizePEMBody(base64Cert)
	der, err := base64.StdEncoding.DecodeString(certBody)
	if err != nil || len(der) == 0 {
		return rsaSHA1
	}

	// First, try to detect SM2 via the certificate's signatureAlgorithm OID
	type algorithmIdentifier struct {
		Algorithm  asn1.ObjectIdentifier
		Parameters asn1.RawValue `asn1:"optional"`
	}
	type certWrapper struct {
		TBS                asn1.RawValue
		SignatureAlgorithm algorithmIdentifier
		SignatureValue     asn1.BitString
	}
	var cw certWrapper
	if _, err := asn1.Unmarshal(der, &cw); err == nil {
		sm2sm3OID := asn1.ObjectIdentifier{1, 2, 156, 10197, 1, 501}
		if cw.SignatureAlgorithm.Algorithm.Equal(sm2sm3OID) {
			return sm2SM3
		}
	}

	// Next, inspect the SubjectPublicKeyInfo algorithm OID for SM2
	cert, err := x509.ParseCertificate(der)
	if err == nil {
		type subjectPublicKeyInfo struct {
			Algorithm        algorithmIdentifier
			SubjectPublicKey asn1.BitString
		}
		var spki subjectPublicKeyInfo
		if _, err := asn1.Unmarshal(cert.RawSubjectPublicKeyInfo, &spki); err == nil {
			sm2PubKeyOID := asn1.ObjectIdentifier{1, 2, 156, 10197, 1, 301}
			if spki.Algorithm.Algorithm.Equal(sm2PubKeyOID) {
				return sm2SM3
			}
		}
		// Fallback to RSA if key algorithm is RSA
		if cert.PublicKeyAlgorithm == x509.RSA {
			return rsaSHA1
		}
	}
	// Default to RSA-SHA1
	return rsaSHA1
}

func (h *ProxyHandler) findCEBClosingTag(xmlContent string) string {
	re := regexp.MustCompile(`ceb:CEB(\d+)Message`)
	m := re.FindStringSubmatch(xmlContent)
	if len(m) > 1 {
		return fmt.Sprintf("</ceb:CEB%sMessage>", m[1])
	}
	return "</ceb:CEB311Message>"
}

func (h *ProxyHandler) insertSignatureIntoXML(originalXML string, signatureXML string) (string, error) {
	closing := h.findCEBClosingTag(originalXML)
	if strings.Contains(originalXML, closing) {
		return strings.Replace(originalXML, closing, signatureXML+closing, 1), nil
	}
	// Fallback: place before the last closing tag in the document
	if i := strings.LastIndex(originalXML, "</"); i != -1 {
		return originalXML[:i] + signatureXML + originalXML[i:], nil
	}
	return "", fmt.Errorf("could not find suitable location to insert signature")
}

func (h *ProxyHandler) processCEBXMLSignature(ctx context.Context, req CEBSignRequest, requestID string, logger *slog.Logger) CEBSignResponse {
	// 1) Canonicalize entire XML using C14N 1.0 (no comments)
	c14n, err := h.canonicalizeXML(req.XML)
	if err != nil {
		logger.Error("Canonicalization failed", "error", err)
		return CEBSignResponse{Success: false, Error: "XML canonicalization failed: " + err.Error()}
	}

	// 2) Compute SHA1 digest (base64) of canonicalized XML for <DigestValue> locally
	digest, err := h.calculateXMLDigest(c14n)
	if err != nil {
		logger.Error("Digest failed", "error", err)
		return CEBSignResponse{Success: false, Error: "XML digest calculation failed: " + err.Error()}
	}

	// 3) Retrieve certificate + key name to populate <KeyInfo>
	certInfo, err := h.getCertificateInfo(ctx, requestID)
	if err != nil {
		logger.Error("Cert info failed", "error", err)
		return CEBSignResponse{Success: false, Error: "Certificate retrieval failed: " + err.Error()}
	}

	// 4) Determine signature algorithm (default rsa-sha1), then build <ds:Signature>
	sigAlgo := determineSignatureMethod(certInfo.Certificate)
	signature := h.createXMLSignature(digest, sigAlgo, certInfo)

	// 5) Canonicalize <SignedInfo> (C14N 1.0, no comments) in the context of the actual document,
	//    then sign using normal API (card performs hashing internally) with the canonicalized SignedInfo text.
	signed, err := h.signXMLSignature(ctx, req.XML, signature, requestID)
	if err != nil {
		logger.Error("Sign XML failed", "error", err)
		return CEBSignResponse{Success: false, Error: "XML signing failed: " + err.Error()}
	}

	logger.Info("CEB XML signed")
	return CEBSignResponse{SignedXML: signed, Success: true}
}

// canonicalizeXML performs W3C Canonical XML 1.0 (no comments) on the **entire document**.
func (h *ProxyHandler) canonicalizeXML(xmlContent string) (string, error) {
	doc := etree.NewDocument()
	if err := doc.ReadFromString(xmlContent); err != nil {
		return "", fmt.Errorf("invalid XML format: %w", err)
	}
	c14n := dsig.MakeC14N10RecCanonicalizer() // inclusive
	out, err := c14n.Canonicalize(doc.Root())
	if err != nil {
		return "", fmt.Errorf("c14n error: %w", err)
	}
	return string(out), nil
}

func (h *ProxyHandler) calculateXMLDigest(canonicalXML string) (string, error) {
	sum := sha1.Sum([]byte(canonicalXML))
	return base64.StdEncoding.EncodeToString(sum[:]), nil
}

func (h *ProxyHandler) getCertificateInfo(ctx context.Context, requestID string) (*CertificateInfo, error) {
	entResp, err := h.ukeyClient.CallWithRequestID(protocol.MethodGetEntName, struct{}{}, requestID)
	if err != nil {
		return nil, fmt.Errorf("failed to get enterprise name: %w", err)
	}
	certResp, err := h.ukeyClient.CallWithRequestID(protocol.MethodGetSignCertAsPEM, struct{}{}, requestID)
	if err != nil {
		return nil, fmt.Errorf("failed to get signing certificate: %w", err)
	}
	noResp, err := h.ukeyClient.CallWithRequestID(protocol.MethodGetCertNo, struct{}{}, requestID)
	if err != nil {
		return nil, fmt.Errorf("failed to get certificate number: %w", err)
	}
	return &CertificateInfo{
		KeyName:       extractStringAt(entResp.Args.Data, 0),
		Certificate:   normalizePEMBody(extractStringAt(certResp.Args.Data, 0)),
		CertificateNo: extractStringAt(noResp.Args.Data, 0),
	}, nil
}

func (h *ProxyHandler) createXMLSignature(digestValue string, signatureMethod string, certInfo *CertificateInfo) *XMLSignature {
	return &XMLSignature{
		XmlnsDS: "http://www.w3.org/2000/09/xmldsig#",
		SignedInfo: SignedInfo{
			CanonicalizationMethod: Method{Algorithm: "http://www.w3.org/TR/2001/REC-xml-c14n-20010315"},
			SignatureMethod:        Method{Algorithm: signatureMethod},
			Reference: Reference{
				URI:        "",
				Transforms: Transforms{Transform: []Method{{Algorithm: "http://www.w3.org/2000/09/xmldsig#enveloped-signature"}}},
				DigestMethod: Method{
					Algorithm: "http://www.w3.org/2000/09/xmldsig#sha1",
				},
				DigestValue: digestValue,
			},
		},
		KeyInfo: KeyInfo{
			KeyName: certInfo.KeyName,
			X509Data: X509Data{
				X509Certificate: certInfo.Certificate,
			},
		},
	}
}

// signXMLSignature inserts a provisional <ds:Signature> (without SignatureValue) into the document,
// canonicalizes <ds:SignedInfo>, asks UKey to sign the canonicalized SignedInfo (card hashes internally),
// sets <ds:SignatureValue>, and finally inserts the full signature before the closing CEB tag.
func (h *ProxyHandler) signXMLSignature(ctx context.Context, originalXML string, signature *XMLSignature, requestID string) (string, error) {
	// Parse original doc
	doc := etree.NewDocument()
	if err := doc.ReadFromString(originalXML); err != nil {
		return "", fmt.Errorf("invalid XML: %w", err)
	}

	// Marshal signature (without SignatureValue) and parse to etree element
	sigBytes, err := xml.Marshal(signature)
	if err != nil {
		return "", fmt.Errorf("failed to marshal provisional Signature: %w", err)
	}
	sigDoc := etree.NewDocument()
	if err := sigDoc.ReadFromBytes(sigBytes); err != nil {
		return "", fmt.Errorf("failed to parse provisional Signature: %w", err)
	}
	sigElem := sigDoc.Root()

	// Append to the root (equivalent to inserting before closing tag)
	root := doc.Root()
	root.AddChild(sigElem)

	// Canonicalize SignedInfo in the context of the document
	c14n := dsig.MakeC14N10RecCanonicalizer() // inclusive, no comments
	signedInfoElem := sigElem.FindElement("./ds:SignedInfo")
	if signedInfoElem == nil {
		// etree's ElementPath matches literal tag names; if prefix handling differs, also try local-name lookup:
		signedInfoElem = sigElem.FindElement("./SignedInfo")
	}
	if signedInfoElem == nil {
		return "", fmt.Errorf("SignedInfo element not found in provisional Signature")
	}
	signedInfoC14N, err := c14n.Canonicalize(signedInfoElem)
	if err != nil {
		return "", fmt.Errorf("failed to canonicalize SignedInfo: %w", err)
	}

	// Ask UKey to sign the canonicalized SignedInfo (normal sign)
	signResp, err := h.ukeyClient.CallWithRequestID(protocol.MethodSignDataAsPEM, map[string]interface{}{
		"inData": string(signedInfoC14N),
	}, requestID)
	if err != nil {
		return "", fmt.Errorf("UKey signing communication error: %w", err)
	}
	if !signResp.Args.Result {
		return "", fmt.Errorf(joinUKeyErrors(signResp.Args.Error))
	}
	signatureValue := extractStringAt(signResp.Args.Data, 0)
	if signatureValue == "" {
		return "", fmt.Errorf("no signature value returned from UKey")
	}
	// Normalize PEM/base64 output to body-only base64 for <ds:SignatureValue>
	signatureValue = normalizePEMBody(signatureValue)

	// Set SignatureValue in the etree node
	sigValueElem := sigElem.FindElement("./ds:SignatureValue")
	if sigValueElem == nil {
		sigValueElem = etree.NewElement("ds:SignatureValue")
		sigElem.AddChild(sigValueElem)
	}
	sigValueElem.SetText(signatureValue)

	// Serialize the signature element only, and inject it into the ORIGINAL XML string before the CEB closing tag.
	// This preserves the original document's prolog/whitespace exactly.
	sigXML := elementToString(sigElem)
	return h.insertSignatureIntoXML(originalXML, sigXML)
}

// elementToString renders an etree.Element without adding XML declaration.
func elementToString(el *etree.Element) string {
	tmpDoc := etree.NewDocument()
	tmpDoc.SetRoot(el.Copy())
	s, err := tmpDoc.WriteToString()
	if err != nil {
		panic(fmt.Sprintf("failed to marshal element to XML: %s", err))
	}
	return strings.TrimSpace(s)
}

// HandleCEBSign handles the /api/ceb/sign endpoint for signing CEB XML documents.
func (h *ProxyHandler) HandleCEBSign(c *gin.Context) {
	requestID, logger, started := h.startRequest(c)

	var req CEBSignRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		logger.Warn("Invalid CEB request", "error", err)
		resp := CEBSignResponse{Success: false, Error: "Invalid request format: " + err.Error()}
		h.sendAudit(c.Request.Context(), requestID, "ceb_xml_sign", nil, resp, false, resp.Error, started)
		c.JSON(http.StatusBadRequest, resp)
		return
	}

	logger.Info("CEB sign request received", "xmlLength", len(req.XML))

	resp := h.processCEBXMLSignature(c.Request.Context(), req, requestID, logger)
	status := httpStatusFromError(resp.Error)

	h.sendAudit(c.Request.Context(), requestID, "ceb_xml_sign", req, resp, resp.Success, resp.Error, started)
	utils.JSONNoEscape(c, status, resp)
}
